package cn.demo.springcloud.filter;

import com.google.common.util.concurrent.RateLimiter;
import com.netflix.zuul.ZuulFilter;
import com.netflix.zuul.context.RequestContext;
import com.netflix.zuul.exception.ZuulException;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpServletRequest;
import java.util.HashMap;
import java.util.Map;

@Component
public class RateLimiterFilter extends ZuulFilter {

    private static Map<String, RateLimiter> IP_RATELIMITER_MAP = new HashMap<>();

    /**
     * 过滤器类型
     * pre：可以在请求被路由之前调用
     * route：在路由请求时候被调用
     * post：在route和error过滤器之后被调用
     * error：处理请求时发生错误时被调用
     */
    @Override
    public String filterType() {
        return "pre";
    }

    /**
     * 过滤器顺序 - 数字小的先执行
     */
    @Override
    public int filterOrder() {
        return 1;
    }

    /**
     * 过滤器是否生效 - 返回true需要过滤，false不需要过滤
     */
    @Override
    public boolean shouldFilter() {
        return true;
    }

    /**
     * 过滤业务 - 只有上面返回true，才会进入该方法
     */
    @Override
    public Object run() throws ZuulException {
        RequestContext requestContext = RequestContext.getCurrentContext();
        HttpServletRequest request = requestContext.getRequest();
        //从request.getHeader("x-forwarded-for")里面获取第一个外网ip，如果找不到则说明是内网访问内网，则取第一个ip做为客户端ip
        //此处客户端ip固定为本地【仅作演示】
        String clientIp = "127.0.0.1";

        if(!IP_RATELIMITER_MAP.containsKey(clientIp)){
            //如果当前客户没有限速，则添加限速【限速每秒1次，仅作演示】
            RateLimiter rateLimiter = RateLimiter.create(1);
            IP_RATELIMITER_MAP.put(clientIp, rateLimiter);
        }
        //获取当前客户限速
        RateLimiter curRateLimiter = IP_RATELIMITER_MAP.get(clientIp);

        if (!curRateLimiter.tryAcquire()) {
            //取不到令牌则超速拒绝
            requestContext.setSendZuulResponse(false);
            requestContext.setResponseStatusCode(HttpStatus.TOO_MANY_REQUESTS.value());
        }
        return null;
    }

}
